from UItools import *
from CustomInteractorStyle import CustomInteractorStyle
from PartData import PartData
from abq_inp_parse import Assembly
from PartDataINP import PartDataINP
from UIconf import *


class VTKContourRenderer(QFrame):
    def __init__(self, parent=None):
        super().__init__(parent)
        
        # 主布局
        self.main_layout = QHBoxLayout()
        self.main_layout.setContentsMargins(0, 0, 0, 0)  # 设置边距为0
        self.main_layout.setSpacing(0)
        self.setLayout(self.main_layout)
        
        # 创建VTK渲染窗口
        self.vtk_widget = QVTKRenderWindowInteractor(self)
        self.main_layout.addWidget(self.vtk_widget)
        
        # 初始化渲染器
        self.renderer = vtk.vtkRenderer()
        self.mesh_visible = False
        self.render_window = self.vtk_widget.GetRenderWindow()
        self.render_window.AddRenderer(self.renderer)
        self.interactor = self.vtk_widget.GetRenderWindow().GetInteractor()
        
        # 初始化颜色映射和标量条
        self.color_map = vtk.vtkLookupTable()
        self._setup_color_map()
        self.scalar_bar = vtk.vtkScalarBarActor()
        self._setup_scalar_bar()
        self.parts = {}
        
        # 设置背景色为白色
        self.renderer.SetBackground(1, 1, 1)

        style = CustomInteractorStyle()
        style.SetDefaultRenderer(self.renderer)  # 关键修复：显式设置渲染器
        self.interactor.SetInteractorStyle(style)
        
        # 完全禁用默认交互
        style.SetAutoAdjustCameraClippingRange(0)
        self.init_view()
        
        # 创建工具栏
        self._create_toolbar()
        
    def _create_toolbar(self):
        """创建底部工具栏"""
        toolbar = QVBoxLayout()
        bkg = QFrame()
        bkg.setStyleSheet("background-color: white;")
        
        # 网格可见性切换按钮
        self.mesh_button = QPushButton("网格")
        self.mesh_button.clicked.connect(self.switch_mesh_visibility)
        self.mesh_button.setStyleSheet(SBUTTON_STYLE)
        toolbar.addWidget(self.mesh_button)
        
        # 重置视图按钮
        self.reset_button = QPushButton("重置")
        self.reset_button.clicked.connect(self.reset_view)
        self.reset_button.setStyleSheet(SBUTTON_STYLE)
        toolbar.addWidget(self.reset_button)
        
        # 添加工具栏到底部
        bkg.setLayout(toolbar)
        self.main_layout.addWidget(bkg)
    
    def add_model(self, model: Assembly):
        """添加装配体"""
        for inst in model.instances:
            part = PartDataINP(inst)
            self.parts[inst.name] = part
            self.renderer.AddActor(part.actor)
            part.only_geom()
            part.set_cmap(self.color_map)
        self._update_global_color_range()
        self.renderer.ResetCamera()

    def load_shared(self, obj: "VTKContourRenderer"):
        """直接引用其他VTK组件数据"""
        for name,part0 in obj.parts.items():
            part = PartData()
            part.load_shared(part0)
            self.parts[name] = part
            self.renderer.AddActor(part.actor)
            part.only_geom()
            part.set_cmap(self.color_map)
        self._update_global_color_range()
        self.renderer.ResetCamera()
    
    def add_part(self, name, node_arr, element_arr, response_arr=None, node_id_arr=None):
        """添加部件并更新全局颜色范围"""
        part = PartData(node_arr, element_arr, node_id_arr)
        self.parts[name] = part
        part.set_cmap(self.color_map)
        self.renderer.AddActor(part.actor)
        if response_arr is not None:
            part.update_contour(response_arr)
        self._update_global_color_range()
        self.renderer.ResetCamera()
    
    def _setup_color_map(self):
        """设置颜色映射表"""
        self.color_map.SetNumberOfColors(256)
        self.color_map.SetHueRange(0.66, 0.0)  # 从蓝到红
        self.color_map.SetSaturationRange(1.0, 1.0)
        self.color_map.SetValueRange(1.0, 1.0)
        self.color_map.SetRange(0, 20)  # 初始范围，加载数据后会更新
        self.color_map.Build()
    
    def _setup_scalar_bar(self):
        """设置标量条(图例)"""
        self.scalar_bar.SetLookupTable(self.color_map)
        self.scalar_bar.SetNumberOfLabels(5)
        self.scalar_bar.SetLabelFormat("%.2f")
        self.scalar_bar.SetPosition(0.9, 0.1)
        self.scalar_bar.SetWidth(0.1)
        self.scalar_bar.SetHeight(0.7)
        # 设置标量条文本颜色为黑色
        self.scalar_bar.GetLabelTextProperty().SetColor(0, 0, 0)
        self.scalar_bar.GetTitleTextProperty().SetColor(0, 0, 0)
        self.scalar_bar.SetVisibility(False)
        
    def enable_contour(self):
        self.scalar_bar.SetVisibility(True)
        self._update_global_color_range()

    def _update_global_color_range(self):
        """联合所有部件的范围更新全局颜色映射"""
        min_vals, max_vals = zip(*[part.get_min_max_response() for part in self.parts.values()])
        global_min, global_max = min(min_vals), max(max_vals)
        print(global_min, global_max)
        self.color_map.SetRange(global_min, global_max)
    
    def set_part_visibility(self, name, visible):
        """控制单个部件可见性"""
        if name in self.parts:
            self.parts[name].set_visibility(visible)
    
    def set_all_edges_visibility(self, visible):
        """控制所有部件的网格边线可见性"""
        for part in self.parts.values():
            part.set_edge_visibility(visible)

    def init_view(self):
        self.renderer.AddActor(self.scalar_bar)
        camera = self.renderer.GetActiveCamera()
        camera.SetPosition(1, 0.5, 1)
        camera.SetFocalPoint(0, 0, 0)
        camera.SetViewUp(0, 1, 0)
    
    def render(self):
        """执行渲染"""
        self.render_window.Render()
        self.vtk_widget.Initialize()
        self.vtk_widget.Start()
    
    def switch_mesh_visibility(self):
        """切换网格可见性"""
        self.mesh_visible = not self.mesh_visible
        self.set_all_edges_visibility(self.mesh_visible)
        self.render_window.Render()
    
    def reset_view(self):
        """重置视图"""
        self.renderer.ResetCamera()
        self.render_window.Render()

    def clear(self):
        """清空渲染内容"""
        self.renderer.RemoveAllViewProps()
        self.render()

    def update(self):
        super().update()
        self.render()
    
    def set_scalar_name(self,name:str):
        self.scalar_bar.SetTitle(name)